#include "OffLineRenderPCH.h"
#include "OR_SamplerStratified.h"

using namespace OffLineRender;
using namespace Utility;
//---------------------------------------------------------------------
OR_SamplerStratified::OR_SamplerStratified(int32_t iXPixelStart, int32_t iXPixelEnd, 
	int32_t iYPixelStart, int32_t iYPixelEnd, 
	int32_t iXPixelSampleCount, int32_t iYPixelSampleCount, 
	bool bJitter, float fShutterOpenTime, float fShutterCloseTime)
	: OR_SamplerBase(iXPixelStart, iXPixelEnd, iYPixelStart, iYPixelEnd,
	iXPixelSampleCount * iYPixelSampleCount, fShutterOpenTime, fShutterCloseTime)
{
	m_bJitterSamples = bJitter;
	m_iXPos = iXPixelStart;
	m_iYPos = iYPixelStart;
	m_iXPixelSampleCount = iXPixelSampleCount;
	m_iYPixelSampleCount = iYPixelSampleCount;

	m_pkSampleBuf = (float*)( MM_STAND_ALLOC(5 * m_iXPixelSampleCount * m_iYPixelSampleCount * sizeof(float)) );
}

//---------------------------------------------------------------------
OR_SamplerStratified::~OR_SamplerStratified()
{
	MM_STAND_DEALLOC(m_pkSampleBuf);
}
//---------------------------------------------------------------------
int OR_SamplerStratified::GenerateSamples(OR_Sample *sample, Utility::UT_RandomNumGenerator &rng)
{
	if (m_iYPos == m_iYPixelEnd)
		return 0;
	
	int nSamples = m_iXPixelSampleCount * m_iYPixelSampleCount;
	// Generate stratified camera samples for _(xPos, yPos)_

	// Generate initial stratified samples into _sampleBuf_ memory
	float *bufp = m_pkSampleBuf;
	float *imageSamples = bufp; bufp += 2 * nSamples;
	float *lensSamples = bufp;  bufp += 2 * nSamples;
	float *timeSamples = bufp;

	UT_SampleHelper::SampleStratified2D(imageSamples, m_iXPixelSampleCount, m_iYPixelSampleCount, rng,m_bJitterSamples);

	UT_SampleHelper::SampleStratified2D(lensSamples, m_iXPixelSampleCount, m_iYPixelSampleCount, rng,m_bJitterSamples);

	UT_SampleHelper::SampleStratified1D(timeSamples, m_iXPixelSampleCount * m_iYPixelSampleCount, rng,m_bJitterSamples);

	// Shift stratified image samples to pixel coordinates
	for (int o = 0; o < 2 * m_iXPixelSampleCount * m_iYPixelSampleCount; o += 2)
	{
		imageSamples[o]   += m_iXPos;
		imageSamples[o+1] += m_iYPos;
	}

	// Decorrelate sample dimensions
	UT_SampleHelper::ShuffleSamples(lensSamples, m_iXPixelSampleCount*m_iYPixelSampleCount, 2, rng);
	UT_SampleHelper::ShuffleSamples(timeSamples, m_iXPixelSampleCount*m_iYPixelSampleCount, 1, rng);

	// Initialize stratified _samples_ with sample values
	for (int i = 0; i < nSamples; ++i)
	{
		sample[i].m_fImageX = imageSamples[2*i];
		sample[i].m_fImageY = imageSamples[2*i+1];
		sample[i].m_fLensU	= lensSamples[2*i];
		sample[i].m_fLensV	= lensSamples[2*i+1];
		sample[i].m_fTime	= UT_MathHelper::Lerp(timeSamples[i], m_fShutterOpenTime, m_fShutterCloseTime);
	
		// Generate stratified samples for integrators
		for (uint32_t j = 0; j < sample[i].m_k1DNumVector.GetSize(); ++j)
			UT_SampleHelper::SampleLatinHypercube(sample[i].m_pp1DData[j], sample[i].m_k1DNumVector.GetAt(j), 1, rng);
		for (uint32_t j = 0; j < sample[i].m_k2DNumVector.GetSize(); ++j)
			UT_SampleHelper::SampleLatinHypercube(sample[i].m_pp2DData[j], sample[i].m_k2DNumVector.GetAt(j), 2, rng);
	}

	// Advance to next pixel for stratified sampling
	if (++m_iXPos == m_iXPixelEnd)
	{
		m_iXPos = m_iXPixelStart;
		++m_iYPos;
	}
	return nSamples;
}
//---------------------------------------------------------------------
OR_SamplerBase* OR_SamplerStratified::GetSubSampler(int num, int count)
{
	int x0, x1, y0, y1;
	_computeSubWindow(num, count, &x0, &x1, &y0, &y1);

	if (x0 == x1 || y0 == y1)
		return NULL;

	return MM_NEW OR_SamplerStratified(x0, x1, y0, y1,
		m_iXPixelSampleCount,m_iYPixelSampleCount, 
		m_bJitterSamples, m_fShutterOpenTime,m_fShutterCloseTime);
}
//---------------------------------------------------------------------
